#
#
import os, sys
import datetime
import tables
import numpy as np
from matplotlib.mlab import find, prctile
from cwd.lib.utils import length, invdist

import pdb
PDB = pdb.set_trace

class PrecipGrid(object):
    def __init__(self, lat, lon, level, filename=None, filepath=None):
        self.lat = lat
        self.lon = lon
        self.level = level

        if filepath is None:
            import cwd.settings as settings
            filepath = settings.DATA_DIR
        if filename is None:
            filename = 'SNOWGRIDMONTH1.h5'
        """
        This file contains two pieces of information in additional to lat/lon and
        a variable called levels.  Levels is an elevational grid from 0 to 4000m
        by 200m.  Unlike the previous mapping exercises that were two dimensional,
        the snow piece is three dimensional since it is elevation dependent.

        The two main variables are PRECIPGRID and data.  PRECIPGRID is the monthly
        cumulative precipitation for each lat/lon/el (fixed over elevation).
        data is the cumulative monthly precipitation that fell on days where the
        freezing level was below the elevation of the grid.
        Hence data<=PRECIPGRID

        The size of these data are:
        20 (lat)   43 (lon)   21 (levels)  12 (month)  65 (years and counting)

        To get % of precip as snow you would just take data./PRECIPGRID*100.
        """

        # load monthly values from big grid, 
        # this is a 5d file dimensions (year,month,levels,lon,lat) = size(num_years, 12, 21, 43, 20)
        filename = os.path.join(filepath,filename)
        hdf = tables.openFile(filename)

        self.precipgrid = self.extract(hdf.root.PRECIPGRID[:])
        self.data = self.extract(hdf.root.data[:])
        self.level_map = hdf.root.levels[:]

        self.last_jday = hdf.root.lastday[0].flatten()[0] # the Julian day of they last data day
        now = datetime.datetime.now()
        self.last_day = datetime.datetime(now.year,1,1) + datetime.timedelta(days=self.last_jday-1)

        #self.p = hdf.root.PRECIPGRID
        #self.d = hdf.root.data
        hdf.close()
        del hdf


    def extract(self, data):
        # transpose to match matlab
        DR = data
        DR = DR.transpose() # now (lat,lon,levels,month,year)
        # Extract the level we want
        DR = np.squeeze(DR[:,:,self.level,:,:])
        DR = np.reshape(DR, (20,43,12*np.size(DR,3)), order='F') # DR=reshape(DR,20,43,12*size(DR,4));
        # inverse distance weight to specified grid
        DATA = invdist(DR,360-self.lon,self.lat)
        DATA = np.reshape(DATA, (12, DATA.size/12), order='F') # DATA=reshape(DATA,12,length(DATA)/12);

        return DATA

    def get_monthly_data(self, month, span=1):
        """
        Return an array of data for a given month, or span of months.
        """
        month = int(month)  # month should be 1-12
        data = self.data
        precipgrid = self.precipgrid
        if span > 1:
            data = data.transpose()
            precipgrid = precipgrid.transpose()

            # Build a list of start/end indicies
            s_i = np.arange(month-span,data.size+1,12)
            e_j = np.arange(month,data.size+1,12)
            # Take the sum over the span and use that as the bar graph data
            d = np.array([data.ravel()[i:j].sum() for i,j in zip(s_i,e_j)])
            p = np.array([precipgrid.ravel()[i:j].sum() for i,j in zip(s_i,e_j)])
            d = d/p
            d *= 100.
            return d
        else:
            data = data/precipgrid
            data *= 100.
            return data[month-1,:]

if __name__ == '__main__':
    lat = 39.2
    lon = 120.2
    level = 10
    d = PrecipGrid(lat,lon,level)
